import time
import numpy as np
import pymysql
import pandas as pd
from datetime import datetime
import pymysql.cursors

class MysqlUtils(object):
    def __init__(self, *args):
        self.conn = pymysql.connect(
            host="127.0.0.1",
            user="root",
            passwd="root",
            db="scenic",
            port=3306,
            charset="utf8"
        )

    def get_scenic_data(self):
        """
        获取数据
        """
        cursor = self.conn.cursor(cursor=pymysql.cursors.DictCursor)
        sql = """
        SELECT LEFT(u.id_no,4) as city_code, DATE_FORMAT(o.create_time, '%Y-%m') as month, COUNT(u.id) as visitor_count
        FROM ticket_order o join ticket_order_user_rel u on o.id = u.order_id
        WHERE LENGTH(u.id_no) =18 and o.pay_time is not null and o.pay_time != ''
        GROUP BY city_code, DATE_FORMAT(o.create_time, '%Y-%m')
        """
        cursor.execute(sql)
        ret = cursor.fetchall()
        new_list = []
        for item in ret:
            if item['city_code'] not in CITY_DICT:
                continue
            new_list.append({
                'city_code': item['city_code'],
                'month': item['month'],
                'visitor_count': item['visitor_count'],
                'city_name': CITY_DICT[item['city_code']]
            })
        df_city_monthly = pd.DataFrame(new_list)
        df_city_monthly['month'] = pd.to_datetime(df_city_monthly['month'] + '-01')
        return df_city_monthly

    def calculate_baseline(df, current_month, window_size=6):
        # 提取当月数据
        df_current = df[df['month'] == current_month]

        # 提取历史数据
        history_start = current_month - pd.DateOffset(month=window_size)
        df_history = df[(df['month'] >= history_start) & (df['month'] < current_month)]

        # 按城市计算均值和标准差
        df_baseline = df_history.groupby('city_name')['visitor_count'].agg(['mean', 'std']).reset_index()
        df_baseline.rename(columns={'mean': 'hist_mean', 'std': 'hist_std'}, inplace=True)

        # 合并当月数据
        df_merged = df_current.merge(df_baseline, on='city_name', how='left')
        return df_merged

# 假设当前月是2024-12
current_month = pd.to_datetime('2024-12-01')
CITY_DICT = {
    '110000': 'Beijing',
    '120000': 'Tianjin',
    '130000': 'Hebei'
}

utils = MysqlUtils()
df_city_monthly = utils.get_scenic_data()
df_merged = utils.calculate_baseline(df_city_monthly, current_month)

# 标记爆发增长的城市（z_score > 3 or z_score < -3）
df_merged['z_score'] = (df_merged['visitor_count'] - df_merged['hist_mean']) / df_merged['hist_std']
df_increased = df_merged[df_merged['z_score'] > 3]
df_reduced = df_merged[df_merged['z_score'] < -3]

print(df_increased)
print('-------------------------')
print(df_reduced)